rusotoでAPIリクエストをタイムアウト付きで実行する

はじめに

オフィシャルのRust AWS SDKが公開されて久しいですがまだrusotoを使っている人は多いと思いいます。今回rusotoに実装されていないけどよく使う2つの機能を自前で実装してみたのでそれぞれ記事にしたいと思います。

やったこと

• 各AWSサービスクライアントのAPIをタイムアウト付きで実行する ← 今回はこっち
• 各AWSサービスのレスポンスに応じてリトライを実行する

作戦

rusoto自体にはリクエストにタイムアウトを設定する機能がないのでtokio::time::timeoutを使ってFutureにタイムアウトを設定します。

タイムアウトを表すためにRusotoError とタイムアウト(Timeout) のどちらかの値を取るenum RusotoErrorWithTimeout を導入します。

タイムアウトの実装例

以下コード例です。タイムアウト時間は固定値にしましたがwith_timeoutで指定できるようにしてもいいと思います。

use rusoto_core::RusotoError;
use std::fmt;
use std::time::Duration;
use std::future::Future;
use tokio::time::timeout;

/// AWSリクエストのタイムアウト 500ms
pub const AWS_REQUEST_TIMEOUT:Duration = Duration::from_millis(500);

/// タイムアウトによる終了を含むRusutoErrorの拡張
#[derive(Debug)]
pub enum RusotoErrorWithTimeout<E> {
  From(RusotoError<E>),
  Timeout,
}

impl <E:fmt::Display+ std::error::Error + 'static> fmt::Display for RusotoErrorWithTimeout<E> {
  fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
      match self {
        RusotoErrorWithTimeout::From(e) => write!(f, "{}", e),
        RusotoErrorWithTimeout::Timeout => write!(f, "timeout after {:}ms", AWS_REQUEST_TIMEOUT.as_millis()),
      }
  }
}

/// タイムアウト付きRusotoResult
pub type TimedRusotoResult<R, E> = Result<R, RusotoErrorWithTimeout<E>>;

/// RusotoのAPIコールFutureをタイムアウト付きに変換する
pub async fn with_timeout<F, R, E>(f:F) -> TimedRusotoResult<R, E>  where F: Future<Output = Result<R, RusotoError<E>>>, {
  match timeout(AWS_REQUEST_TIMEOUT, f).await {
      Ok(r) => r.map_err(RusotoErrorWithTimeout::From),
      Err(_) => Err(RusotoErrorWithTimeout::Timeout),
  }
}

APIクライアントの実装例

実際に使うときには以下のようにrusoto_dynamodbをラップしました。

use rusoto_dynamodb::{
  DynamoDb,
  DynamoDbClient as Inner, 
  PutItemInput, PutItemOutput,  PutItemError
};
pub struct DynamoDbClient {
  inner: Inner,
}
pub async fn put_item(
    &self,
    input: PutItemInput
  ) -> TimedRusotoResult<PutItemOutput, PutItemError>{
    with_timeout(self.inner.put_item(input)).await
  }

rusotoにおけるタイムアウトの実装

coreで定義されているHTTPクライアントレベルのtrait SignAndDispatchにはタイムアウトの指定ができてインスタンスに実装されているのですが、これを使っているAWS APIクライアントレベルではタイムアウトは設定していないので利用できません。

localstackでのタイムアウトのテスト

rusotoに限らずAWSサービスを使った開発ではローカルでのスタブとしてlocalstackを使うことが多いと思います。tcコマンドを使ってネットワークのレイテンシを追加することでタイムアウトが発生する状況を再現できます。

参考: tcコマンドとDockerコンテナを用いて遅いネットワークをシミュレートする

まとめ

tokio::time::timeoutを使ってrusotoでタイムアウト付きのAPIリクエストを実装しました。